from .base import BaseDataset
import os

class Cub(BaseDataset):
    nb_train_all = 5864
    nb_test_all = 5924
    def __init__(self, root, classes, transform=None):
        BaseDataset.__init__(self, root, classes, transform)

        # classes_train = range(0, 100)
        # classes_test = range(100, 200)

        if classes == range(0, 100):
            train = True
        else:
            train = False
        # if classes.start in classes_train:
        #     if classes.stop - 1 in classes_train:
        #         train = True
        #
        # if classes.start in classes_test:
        #     if classes.stop - 1 in classes_test:
        #         train = False

        with open(
            os.path.join(
                root,
                '{}100.txt'.format('train' if train else 'test')
            )
        ) as f:

            # f.readline()
            index = 0
            nb_images = 0

            for (_, path) in map(str.split, f):
                class_id = int(path.split('.')[0])
                nb_images += 1
                self.im_paths.append(os.path.join(root, 'images', path))
                self.ys.append(class_id-1)
                self.I += [index]
                index += 1

            if train:
                assert nb_images == type(self).nb_train_all
            else:
                assert nb_images == type(self).nb_test_all
